import torch
from torch.utils.data import DataLoader
from utils import get_datasets
from utils import label_to_onehot, onehot_to_label

DATA_PATH = './data/'
train_dataset, test_dataset = get_datasets(DATA_PATH)
train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=True)

print('train_dataset size: ', len(train_dataset))
print('test_dataset size: ', len(test_dataset))

print(onehot_to_label(label_to_onehot('mbfd')))
print(onehot_to_label(label_to_onehot('abcd')))

for X, y in train_dataset:
    print(f'y_label: {onehot_to_label(y)}')

for X, y in train_dataloader:
    print(f'X shape: {X.shape}, y shape: {y.shape}')
    for i in range(y.size(0)):
        print(f'y_label: {onehot_to_label(y[i])}')